# https://pytorch.org/tutorials/intermediate/ddp_tutorial.html

import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from PyCmpltrtok.common import sep

# On Windows platform, the torch.distributed package only
# supports Gloo backend, FileStore and TcpStore.
# For FileStore, set init_method parameter in init_process_group
# to a local file. Example as follow:
# init_method="file:///f:/libtmp/some_file"
# dist.init_process_group(
#    "gloo",
#    rank=rank,
#    init_method=init_method,
#    world_size=world_size)
# For TcpStore, same way as on Linux.


def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    # dist.init_process_group("gloo", rank=rank, world_size=world_size)
    dist.init_process_group("nccl", rank=rank, world_size=world_size)


def cleanup():
    dist.destroy_process_group()


class ToyMpModel(nn.Module):
    def __init__(self, dev0, dev1):
        super(ToyMpModel, self).__init__()
        self.dev0 = dev0
        self.dev1 = dev1
        self.net1 = torch.nn.Linear(10, 10).to(dev0)
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5).to(dev1)

    def forward(self, x):
        print('x.to', self.dev0)
        x = x.to(self.dev0)
        print('x.to', self.dev0, 'done')
        print('net1')
        x = self.relu(self.net1(x))
        print('net1 done')
        print('x.to', self.dev1)
        x = x.to(self.dev1)
        print('x.to', self.dev1, 'done')
        print('net2')
        x = self.net2(x)
        print('net2 done')
        return x


def demo_model_parallel(rank, world_size):
    flag = f'{rank}/{world_size}'
    sep(f'{flag} start')
    print(f"Running DDP with model parallel example on rank {rank}.")
    setup(rank, world_size)

    # setup mp_model and devices for this process
    print(flag, 'init model')
    # dev0 = (rank * 2) % world_size
    dev0 = (rank * 2)
    # dev1 = (rank * 2 + 1) % world_size
    dev1 = (rank * 2 + 1)
    mp_model = ToyMpModel(dev0, dev1)
    ddp_mp_model = DDP(mp_model)
    print(flag, 'init model done')

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_mp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    # outputs will be on dev1
    print(flag, 'inferring ...')
    outputs = ddp_mp_model(torch.randn(20, 10))
    print(flag, '20x10 ->', outputs.shape)
    labels = torch.randn(20, 5).to(dev1)
    loss_fn(outputs, labels).backward()
    optimizer.step()

    print(f'{flag} start cleanup')
    cleanup()
    print(f'{flag} cleanup end')
    sep(f'{flag} over')


def run_demo(demo_fn, world_size):
    sep('spawn started (joined)')
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             join=True)
    sep('spawn over')


if '__main__' == __name__:

    # run_demo(demo_model_parallel, 2)
    run_demo(demo_model_parallel, 1)
